package io.gitee.luoyuan800;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestClient;

import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Service
public class SerpApiMCPServer {
    private static final Logger logger = LoggerFactory.getLogger(SerpApiMCPServer.class);
    private final RestClient restClient;
    private final String apiKey;

    public SerpApiMCPServer(@Value("${serpapi.apikey}") String apiKey) {
        this.apiKey = apiKey;
        this.restClient = RestClient.builder()
                .baseUrl("https://serpapi.com")
                .defaultHeader("Accept", "application/json")
                .build();
    }

    @Tool(description = "Google image search, return list of original image urls")
    public List<Map<String, String>> googleImageSearch(
            @ToolParam(description = "Search keyword")
            String keyWord,
            @ToolParam(required = false, description = "Where you want the search to originate")
            String myLocation,
            @ToolParam(required = false, description = "Time period to search for the recent images (d=day, w=week, m=month, y=year)")
            String periodUnit,
            @ToolParam(required = false, description = "An optional time period value")
            Integer periodValue
    ) {
        // Build request URL with parameters
        StringBuilder urlBuilder = new StringBuilder("/search?engine=google_images&q=");
        urlBuilder.append("\"").append(keyWord).append("\"");

        if (myLocation != null && !myLocation.isEmpty()) {
            urlBuilder.append("&location=");
            urlBuilder.append(URLEncoder.encode(myLocation, StandardCharsets.UTF_8));
        }

        if (periodUnit != null && !periodUnit.isEmpty() && periodValue != null) {
            urlBuilder.append("&periodUnit=");
            urlBuilder.append(URLEncoder.encode(periodUnit, StandardCharsets.UTF_8));
            urlBuilder.append("&periodValue=");
            urlBuilder.append(periodValue);
        }

        // Add Google domain parameter for better results
        urlBuilder.append("&google_domain=google.com");
        // Add key
        urlBuilder.append("&api_key=" + apiKey);

        String requestUrl = urlBuilder.toString();
        logger.info("Sending request to SerpApi: {}", requestUrl);

        // Send request to SerpApi
        ResponseEntity<String> response = restClient.get()
                .uri(requestUrl)
                .retrieve()
                .toEntity(String.class);

        logger.info("Received response with status code: {}", response.getStatusCode());

        if (!response.getStatusCode().is2xxSuccessful()) {
            logger.error("Request failed with status code: {}", response.getStatusCode());
            return List.of(Map.of("Error", "Failed to fetch image search results, status code: " + response.getStatusCode()));
        }

        try {
            logger.debug("Parsing response body");
            // Parse JSON response
            JsonNode rootNode = new ObjectMapper().readTree(response.getBody());
            JsonNode imagesResultsNode = rootNode.path("images_results");

            // Extract original image URLs
            List<Map<String, String>> imageUrls = new ArrayList<>();
            if (imagesResultsNode.isArray()) {
                for (JsonNode result : imagesResultsNode) {
                    String originalUrl = result.path("original").asText();
                    String title = result.path("title").asText();
                    if (!originalUrl.isEmpty()) {
                        imageUrls.add(Map.of(title, originalUrl));
                    }
                }
            }

            logger.info("Successfully processed image search, found {} results", imageUrls.size());
            return imageUrls;
        } catch (Exception e) {
            logger.error("Error processing image search results: {}", e.getMessage(), e);
            return List.of(Map.of("Error", "Failed to process image search results: " + e.getMessage()));
        }
    }

    @Tool(description = "Google web search, return organic results")
    public List<Map<String, Object>> googleSearch(
            @ToolParam(description = "Search keyword")
            String keyWord,
            @ToolParam(required = false, description = "Where you want the search to originate")
            String myLocation,
            @ToolParam(required = false, description = "Time period to search for the recent results (d=day, w=week, m=month, y=year)")
            String periodUnit,
            @ToolParam(required = false, description = "An optional time period value")
            Integer periodValue
    ) {
        // Build request URL with parameters
        StringBuilder urlBuilder = new StringBuilder("/search?engine=google_light&q=");
        urlBuilder.append("\"").append(keyWord).append("\"");

        if (myLocation != null && !myLocation.isEmpty()) {
            urlBuilder.append("&location=");
            urlBuilder.append(URLEncoder.encode(myLocation, StandardCharsets.UTF_8));
        }

        if (periodUnit != null && !periodUnit.isEmpty() && periodValue != null) {
            urlBuilder.append("&periodUnit=");
            urlBuilder.append(URLEncoder.encode(periodUnit, StandardCharsets.UTF_8));
            urlBuilder.append("&periodValue=");
            urlBuilder.append(periodValue);
        }

        // Add Google domain parameter for better results
        urlBuilder.append("&google_domain=google.com");
        // Add key
        urlBuilder.append("&api_key=" + apiKey);
        String requestUrl = urlBuilder.toString();
        logger.info("Sending request to SerpApi: {}", requestUrl);

        // Send request to SerpApi
        ResponseEntity<String> response = restClient.get()
                .uri(requestUrl)
                .retrieve()
                .toEntity(String.class);

        logger.info("Received response with status code: {}", response.getStatusCode());

        if (!response.getStatusCode().is2xxSuccessful()) {
            logger.error("Request failed with status code: {}", response.getStatusCode());
            return List.of(Map.of("Error", "Failed to fetch search results, status code: " + response.getStatusCode()));
        }

        try {
            // Parse JSON response
            JsonNode rootNode = new ObjectMapper().readTree(response.getBody());
            JsonNode organicResultsNode = rootNode.path("organic_results");

            // Extract organic results
            List<Map<String, Object>> organicResults = new ArrayList<>();
            if (organicResultsNode.isArray()) {
                for (JsonNode result : organicResultsNode) {
                    Map<String, Object> resultMap = new java.util.HashMap<>();
                    resultMap.put("position", result.path("position").asInt());
                    resultMap.put("title", result.path("title").asText());
                    resultMap.put("link", result.path("link").asText());
                    resultMap.put("displayed_link", result.path("displayed_link").asText());
                    resultMap.put("snippet", result.path("snippet").asText());

                    // Add sitelinks if present
                    JsonNode sitelinksNode = result.path("sitelinks").path("inline");
                    if (sitelinksNode.isArray() && sitelinksNode.size() > 0) {
                        List<Map<String, String>> sitelinks = new ArrayList<>();
                        for (JsonNode sitelink : sitelinksNode) {
                            sitelinks.add(Map.of(
                                    "title", sitelink.path("title").asText(),
                                    "link", sitelink.path("link").asText()
                            ));
                        }
                        resultMap.put("sitelinks", sitelinks);
                    }

                    organicResults.add(resultMap);
                }
            }

            logger.info("Successfully processed web search, found {} organic results", organicResults.size());
            return organicResults;
        } catch (Exception e) {
            logger.error("Error processing web search results: {}", e.getMessage(), e);
            return List.of(Map.of("Error", "Failed to process search results: " + e.getMessage()));
        }
    }
}
